"""
-*- coding: utf-8 -*-
@Author : yyc
@Time : 2023/7/27 17:57
@File : conn_milvus.py
"""
# -*- coding: utf-8 -*-
#!/usr/local/bin/python
# Author: yuxiang yu

'''
国产向量数据库

安装依赖
pip install pymilvus==2.2.13

milvus 版本
2.2.11

# sort
1. create collection
2. crete index
3. load
4. search
'''

from pymilvus import connections, utility, FieldSchema, CollectionSchema, DataType, Collection, Partition
from logging import getLogger
from django.conf import settings
import datetime
logger = getLogger("log")

class MilvusDB(object):
    global _ALIAS
    global _HOST
    global _PORT
    global _USER
    global _PASSWORD
    global _PSW
    global _INDEX_FILE_SIZE
    global _NLIST
    global _NPROBE
    _ALIAS = "default"
    # _HOST = "localhost"
    # _PORT = "19530"

    # 开发环境
    _HOST = settings.MILVUS_HOST
    _PORT = settings.MILVUS_PORT
    _USER = settings.MILVUS_USER
    _PASSWORD = settings.MILVUS_PASSWORD

    # _HOST = "106.75.215.89"
    # _PORT = 24006
    # _USER = "root"
    # _PSW = "Huayuan@2022"
    # MILVUS_SECURE = False
    #
    # MILVUS_COLLECTION_NAME = 'langchain'

    _INDEX_FILE_SIZE = 1024
    _NLIST = 2048
    _NPROBE = 16

    def __init__(self, dimension=768, collection_name="data", output_fields=None, only_search=False, index_type='IP'):
        self.rq = datetime.datetime.now()
        self._DIM = dimension
        self.collection_name = collection_name
        assert index_type in ['L2', 'IP'], 'index type must be L2 or IP'
        # if index_type == 'L2':
        #     # self.search_params = {
        #     #         "metric_type": "L2",
        #     #         "index_type": "IVF_FLAT",
        #     #         "params": {"nlist": 128},
        #     #     }
        #     self.search_params = {
        #             "metric_type": "L2",
        #             "index_type": "FLAT"
        #         }
        # elif index_type == 'IP':
        #     # self.search_params = {
        #     #         "metric_type": "IP",
        #     #         "index_type": "IVF_FLAT",
        #     #         "params": {"nlist": 128},
        #     #     }
        #     self.search_params = {
        #             "metric_type": "IP",
        #             "index_type": "FLAT"
        #         }

        if index_type == 'L2':
            self.search_params = {
                "metric_type": "L2",
                "index_type": "HNSW",
                "params": {"M": 64, "efConstruction": 300}
            }
        elif index_type == 'IP':
            self.search_params = {
                "metric_type": "IP",
                "index_type": "HNSW",
                "params": {"M": 64, "efConstruction": 300}
            }
        assert self._conn() is True
        self.load_collection()
        if only_search:# only search
            if not self.collection.has_index():
                self.set_index()
            self.collection.load()
        if output_fields is None:
            self.output_fields = ['uid', 'text', 'source', 'meta_data', 'embedding', 'project_name', 'base_name']
        else:
            self.output_fields = output_fields

    def load_collection(self):
        try:
            self.collection = self.load_or_create_collection()
            # self.collection.load()
        except:
            self.load_or_create_collection().release()
            self.collection = self.load_or_create_collection()
            # self.collection.load()

    def __del__(self):
        try:
            # self.collection.release()
            self._disconnect()
        except Exception as e:
            print(f"auto release failed {e}")

    def __repr__(self):
        return f"db [{_HOST}], rq {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"

    def _conn(self):
        if _USER and _PASSWORD:
            connections.connect(
                alias=_ALIAS,
                host=_HOST,
                port=_PORT,
                user=_USER,
                password=_PASSWORD,
                db_name=settings.MILVUS_DB_NAME,  # 默认数据库
            )
        else:
            connections.connect(
                alias=_ALIAS,
                host=_HOST,
                port=_PORT,
                db_name=settings.MILVUS_DB_NAME,  # 默认数据库
            )
        try:
            if connections.has_connection(_ALIAS):
                print("connection success")
                return True
            return False
        except Exception as e:
            print(f"connection error: {e}")
            exit(-1)
            return False

    def _disconnect(self):
        try:
            connections.disconnect(_ALIAS)
        except Exception as e:
            print(f"disconnect error: {e}")

    def manual_close(self):
        try:
            self.collection.release()
            print("collection release success")
            connections.disconnect(_ALIAS)
            print("connection close success")
        except Exception as e:
            print(f"disconnect error: {e}")

    @property
    def show_all_columns(self):
        return ["uid", "text", "source", "meta_data", "feature", "embedding", "project_name", "base_name", "create_time"]

    def set_schema(self):
        # uid, embeddings, text,
        fields_list = [
            FieldSchema(name="uid", dtype=DataType.VARCHAR, is_primary=True, auto_id=False, max_length=64),
            FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=3000),
            FieldSchema(name="source", dtype=DataType.VARCHAR, max_length=500),
            FieldSchema(name="meta_data", dtype=DataType.VARCHAR, max_length=3000),
            # FieldSchema(name="label", dtype=DataType.VARCHAR, max_length=100),
            FieldSchema(name="feature", dtype=DataType.FLOAT_VECTOR, dim=self._DIM),
            FieldSchema(name="embedding", dtype=DataType.VARCHAR, max_length=64),
            FieldSchema(name="project_name", dtype=DataType.VARCHAR, max_length=64),
            FieldSchema(name="base_name", dtype=DataType.VARCHAR,  max_length=64),
            FieldSchema(name="create_time", dtype=DataType.VARCHAR, max_length=50)
        ]
        schema = CollectionSchema(
            fields=fields_list,
            description=f"collection: {self.collection_name}",
            enable_dynamic_field=True
            )
        return schema

    def load_or_create_collection(self):
        # TODO: collection whether or not it exists
        if utility.has_collection(self.collection_name, using=_ALIAS):
            collection = Collection(self.collection_name, using=_ALIAS)
            # collection.load(replica_number=2)
            return collection

        # TODO: create a collection with the schema
        try:
            print(f"init databse, start create collection")
            schema = self.set_schema()
            collection = Collection(
                    name=self.collection_name,
                    schema=schema,
                    using=_ALIAS,
                    # consistency_level="Strong",
                    shards_num=2,
                )
            # create index
            self.set_index()
        except:
            pass
        return collection

    def modify_collection(self):
        # todo: old method
        # 生存时间验证
        self.collection.set_properties(properties={"collection.ttl.seconds": 1800})
        return True

    def check_collection(self):
        # 检查 collection
        collection = self.collection
        desc = f"""schema: {collection.schema}\ndescription: {collection.description}\nname: {collection.name}\nis_empty: {collection.is_empty}\nnum_entities: {collection.num_entities}\nprimary_field: {collection.primary_field}\npartitions: {collection.partitions}\nindexes: {collection.indexes}\n"""
        print(desc)

    def get_num_entities(self):
        return self.collection.num_entities

    def get_primary_field(self):
        return self.collection.primary_field

    def get_partitions(self):
        return self.collection.partitions

    def list_collection(self):
        try:
            # self.collection.load()
            print(self.collection)
            return utility.list_collections()
        except Exception as e:
            print(f"list_collection: {e}")
            return None

    def list_partition(self):
        try:
            return self.collection.partitions
        except Exception as e:
            print(f"list partition error: {e}")
            return False

    def drop_collection(self, collection_name=None):
        if collection_name is None:
            collection_name = "models"

        try:
            utility.drop_collection(self.collection_name)
            print(f"drop collection success")
            return True
        except Exception as e:
            print(f"drop collection error: {e}")
            return False

    def drop_partition(self, partition_name):
        try:
            self.collection.drop_partition(partition_name)
            return True
        except Exception as e:
            print(f"drop partition error: {e}")
            self.collection.release()
            self.collection.drop_partition(partition_name)
            return False

    def load_partition(self, partition_name):
        if isinstance(partition_name, str):
            partition_name = [str(partition_name).strip()]

        if not isinstance(partition_name, list):
            partition_name = [partition_name]

        try:
            self.collection.load(partition_name, replica_number=2)
            return True
        except Exception as e:
            print(f"load partition error: {e}")
            return False

    def release_collection(self):
        try:
            self.collection.release()
            return True
        except Exception as e:
            print(f"release collection error: {e}")
            return False

    def release_partition(self, partition_name):
        if isinstance(partition_name, str):
            partition_name = str(partition_name).strip()
            partition = Partition(self.collection, partition_name)
            partition.release()
            return True

        if isinstance(partition_name, (tuple, list)):
            for name in partition_name:
                try:
                    partition = Partition(self.collection, name)
                    partition.release()
                except Exception as e:
                    print(f"release partition error: {e}")
            return True
        print(f"partition_name: {partition_name}")
        return False

    def delete_entity(self, expr, partition_name=None):
        # expr = "sku in [0,1]"
        try:
            self.collection.delete(expr=expr,partition_name=partition_name)
            return True
        except Exception as e:
            print(f"delete entity error: {e}")
            return False

    def set_partition(self, partition_name=None, timeout=20):
        # TODO: create a partition
        if not self.collection.has_partition(
                partition_name=partition_name,
                timeout=timeout):
            try:
                if partition_name is None or str(partition_name).startswith("None"):
                    partition_name = "part1"
                self.collection.create_partition(
                    partition_name=partition_name,
                    description=partition_name
                    )
                return True
            except Exception as e:
                print(f"set partition error:{e}, rq: {self.rq}")
        print(f"set_partition: {partition_name} 已存在, rq:{self.rq}")
        return True

    def set_index(self):
        # if metric_type == 'L2':
        #     index_params = {
        #             "metric_type": "L2",
        #             "index_type": "IVF_FLAT",
        #             "params": {"nlist": 128},
        #         }
        # else:
        #     index_params = {
        #             "metric_type": "IP",
        #             "index_type": "IVF_FLAT",
        #             "params": {"nlist": 128},
        #         }
        if not self.collection.has_index():
            try:
                self.collection.create_index(field_name="feature", index_params=self.search_params)
                # utility.index_building_progress(self)
                print("build or rebuild index success")
                # TODO: reload
                self.collection.load()
                return True
            except Exception as e:
                print(f"build index error: {e}")
                return False
        return True

    def drop_index(self):
        # TODO: the fuction removes all corresponding index files
        try:
            self.collection.drop_index()
            return True
        except Exception as e:
            print(f"drop index error: {e}")
            return False

    def delete_data(self, data, partition_name: str):
        # 首列唯一索引
        # print(self.check_collection())
        idx_list = data[0]
        if not isinstance(idx_list, list):
            print("idx list must be type of list")
            exit(-1)
        expr = f"uid in {str(idx_list)}"
        print(expr)
        try:
            self.collection.delete(expr, partition_name=partition_name)
            print("Remove duplicates finished, delete success!")
            return True
        except Exception as e:
            
            print(f"delete data error {e}")
            return False

    def insert_data(self, data, partition_name: str, timeout=30):
        try:
            if not connections.has_connection(_ALIAS):
                self._conn()
                self.load_collection()
                if not self.collection.has_index():
                    self.set_index()
            # self._conn()
            self.set_partition(partition_name=partition_name, timeout=timeout)
        except Exception as e:
            print(f"set_partition error: {e}")
        init_insert_status = None
        # if not self.collection.has_index():
        #     self.set_index()
        # self.collection.load()
        # init_insert_status = self.collection.insert(data, partition_name=partition_name)
        # # TODO: rebuild index
        # utility.wait_for_index_building_complete(self.collection_name)
        # return True, "success"
        try:
            if not self.collection.has_index():
                self.set_index()
            # self.collection.load()
            init_insert_status = self.collection.insert(data, partition_name=partition_name)
            # TODO: rebuild index
            utility.wait_for_index_building_complete(self.collection_name)
            return True, "success"
        except Exception as e:
            if not self.collection.has_index():
                self.set_index()
            print(f"[insert error] {e} [partition_name] {partition_name} with data {data[0]} ")
            self.collection.release()
            self.collection.load()
            init_insert_status = self.collection.insert(data, partition_name=partition_name)
            return True, "success"

    def _result_dict(self, hit):
        tmp_dict = {
            "uid": hit.entity.get("uid"),
            "text": hit.entity.get("text"),
            # "label": hit.entity.get("label"),
            "dis": hit.distance,
            }
        return tmp_dict

    def get_search(self, query_vector, partition_names, topk=10, timeout=20, expr=None, is_dict=False, search_params=None):
        if not search_params:
            search_params = {
                "metric_type": "IP",
                "params": {"ef": 100},
            }

        if not isinstance(partition_names, list):
            partition_names = [partition_names]

        # if dis-connection try conn
        if not connections.has_connection(_ALIAS):
            self._conn()
            self.load_collection()
            if not self.collection.has_index():
                self.set_index()
            self.collection.load()
        valid_partition_list = []
        for p in partition_names:
            if self.collection.has_partition(
                    partition_name=p,
                    timeout=timeout
            ):
                valid_partition_list.append(p)
        if not valid_partition_list:
            return {0: []}
        try:
            result = self.collection.search(
                data=query_vector,
                anns_field="feature",
                param=search_params,
                limit=topk,
                expr=expr,
                output_fields=self.output_fields,
                timeout=timeout,
                consistency_level="Strong",
                partition_names=valid_partition_list
                )
            if is_dict:
                res_dict = {}
                for idx, hits in enumerate(result):
                    res_dict[idx] = list(map(self._result_dict, list(hits)))
                return res_dict
        except Exception as e:
            print(f"get_search: {e}")
            self.collection.load()
            result = self.collection.search(
                data=query_vector,
                anns_field="feature",
                param=search_params,
                limit=topk,
                expr=expr,
                output_fields=self.output_fields,
                timeout=timeout,
                consistency_level="Strong",
                partition_names=valid_partition_list
                )
            if is_dict:
                res_dict = {}
                for idx, hits in enumerate(result):
                    res_dict[idx] = list(map(self._result_dict, list(hits)))
                return res_dict
        return result


    def get_hybrid_vector_search(self, query_vector, nprobe=10, limit=2, item='2', is_dict=False):
        search_param = {
            "data": query_vector,
            "anns_field": "feature",
            "param": {"metric_type": "IP", "params": {"nprobe": nprobe}},
            "offset": 0,
            "limit": int(limit),
            "expr": f"item == '{item}'",
        }
        # try:
        #     self.set_index()
        # except:
        #     pass
        result = self.collection.search(**search_param)

        return result

    def filter_by_expr(self, expr: str):
        # if dis-connection try conn
        if not connections.has_connection(_ALIAS):
            self._conn()
            self.load_collection()
            if not self.collection.has_index():
                self.set_index()
            self.collection.load()

        try:
            self.collection.load()
        except Exception as e:
            pass
        res = self.collection.query(
            expr=expr,
            offset=0,
            limit=1,
            output_fields=self.output_fields,
        )
        return res


if __name__ == "__main__":
    # 构建一批模型数据
    from uuid import uuid4
    import numpy as np
    import datetime
    nums = 10

    # db = MilvusDB(collection_name='knowledge_test')
    # 删除 collect
    # db.drop_collection()
    # 显示所有列
    # print(db.show_all_columns)
    # 模拟数据
    # TODO: 写入数据
    db = MilvusDB(collection_name='custom_qa_768_IP')
    print(db.check_collection())
    # print(db.collection.partitions)
#     tmp_result_list = []
#     no = [uuid4().__str__() + "-" + datetime.datetime.now().strftime('%Y%m%d%H%M%S') for _ in range(nums)]
#     path = ["/data/pic/" + uuid4().__str__() + ".png" for _ in range(nums)]
#     sku = [uuid4().__str__() for _ in range(nums)]
#     label = [str(np.random.randint(50)) for _ in range(nums)]
#     # item = [str(np.random.randint(50)) for _ in range(nums)]
#     feature = [np.random.random((1, 768)).tolist()[0] for _ in range(nums)]
#     info = [str(np.random.randint(50)) for _ in range(nums)]
#     status = [str(1) for _ in range(nums)]
#     create_time = [datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S') for _ in range(nums)]
#     tmp_result_list.extend([no, path, sku, item, feature, info, status, create_time])
#     # 写入数据
#     # db.insert_data(data=tmp_result_list, partition_name='test')
#     # 手动关闭
#     # db.manual_close()
#     # del db

#     # TODO: 检索数据
#     # 检索
#     db_search = MilvusDB(collection_name='knowledge_test', only_search=True)
#     result = db_search.get_search([feature[0]], partition_names=['test', '_default'], topk=10, nprobe=10, is_dict=True)
#     print(f"识别结果", result)

#     result = db_search.get_search([feature[1]], partition_names=['test', '_default'], topk=10, nprobe=10, is_dict=True)
#     print(f"识别结果", result)
#     # 手动释放
#     db_search.manual_close()
#     del db_search

#     # TODO: 其他指标
#     print("=== 其他指标 ===")
#     db_info = MilvusDB()
#     print(db_info.get_num_entities())
#     db_info.manual_close()
#     del db_info
